# #!/usr/bin/env python

from roboticstoolbox import DHRobot, RevoluteDH, PrismaticDH
from math import pi
import numpy as np
from spatialmath import SE3
import time


class Project_Arm(DHRobot):

    def __init__(self):

        deg = pi / 180
        mm = 1
        
        # P1 : d1初值为(129.22+360*3/17)*mm
        P1 = PrismaticDH(
            # theta=pi/2,
            theta=0,
            a=0.,
            alpha=-pi/2,
            I=[0.0032, 0.0070, 0.0046, 1.178e-06, -6.786e-05, -3.978e-06],  # inertia tensor
            r=[-0.1213, -0.0450, -0.2481],  # center of mass
            m=2.4461,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-0.1, 0.1],  # joint limits (placeholder)
        )

        R1 = RevoluteDH(
            # d=128.5*mm,
            d=0*mm,
            a=0.,
            alpha=-pi/2,
            I=[0.0015, 0.0016, 0.0025, 3.232e-06, 1.444e-06, 5.134e-05],  # inertia tensor
            r=[-0.0093, 0.0001, -0.0301],  # center of mass
            m=1.2270,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
        )

        R2 = RevoluteDH(
            d=0.,
            a=325.9*mm,
            alpha=pi,
            I=[0.0005, 0.0002, 0.0005, -5.364e-09, 1.538e-05, 9.547e-07],  # inertia tensor
            r=[0.0025, -0.1921, 0.0229],  # center of mass
            m=0.5588,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
        )

        R3 = RevoluteDH(
            d=59.6*mm,
            a=239.78*mm,
            alpha=pi,
            I=[0.0001, 0.0001, 0.0002, -5.340e-09, 4.701e-06, -1.946e-07],  # inertia tensor
            r=[-0.0002, 0.0028, -0.0188],  # center of mass
            m=0.3504,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
        )
        
        # R4: theta4的初值为 pi/2
        R4 = RevoluteDH(
            d=59.53*mm,
            a=0.,
            alpha=pi/2,
            I=[4.625e-05, 7.494e-05, 4.471e-05, 1.662e-07, -9.595e-08, -1.336e-06],  # inertia tensor
            r=[-0.0006, 0.0317, 0.0003],  # center of mass
            m=0.1831,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
            offset=pi/2
        )

        # R5: theta5的初值为 pi/2
        R5 = RevoluteDH(
            d=66.14*mm,
            a=0.,
            alpha=pi/2,
            I=[0.0001, 0.0001, 0.0002, -5.340e-09, 4.701e-06, -1.946e-07],  # inertia tensor
            r=[-0.0002, 0.0028, -0.0188],  # center of mass
            m=0.3504,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
            offset=-pi/2
        )

        R6 = RevoluteDH(
            d=104.5*mm,
            a=0.,
            alpha=0,
            I=[4.625e-05, 7.494e-05, 4.471e-05, 1.662e-07, -9.595e-08, -1.336e-06],  # inertia tensor
            r=[-0.0006, 0.0317, 0.0003],  # center of mass
            m=0.1831,  # mass of link
            Jm=0.001,  # actuator inertia (placeholder)
            G=1,  # gear ratio
            qlim=[-170 * deg, 170 * deg],  # joint limits
        )

        L = [P1,R1,R2,R3,R4,R5,R6]

        super().__init__(
            L,
            name="Project arm",
            manufacturer="Victor Scheinman",
            keywords=("dynamics",),
        )

        self.qr = np.zeros(7)
        self.qz = np.zeros(7)

        self.addconfiguration("qr", self.qr)
        self.addconfiguration("qz", self.qz)

if __name__ == "__main__":
    arm = Project_Arm()
    
    #  650.043213,27.257435,-14.799980,-0.909512,-0.745723,-0.835655 
    target = np.array([ 650.043213,27.257435,-14.799980,-0.909512,-0.745723,-0.835655 ])
    T = SE3(target[0],target[1],target[2]) * SE3.RPY(target[3],target[4],target[5])
    q = arm.ikine_LM(T,joint_limits=True)
    
    print("Target pose:")
    print(T)
    print("Joint angles:")
    print(q.q)
    
    # Plot the robot with pyplot
    arm.plot(q.q, block=True)


if __name__ == "__main__1":
    arm=Project_Arm()
    ang = np.array([0,0.1,0.2,0.3,0.4,0.5,0.6])
    T=arm.fkine(ang)
    print("Position (x,y,z):", T.t)  # 打印位置坐标
    print("Euler angles (ZYX convention) in radians:", T.eul())  # 打印欧拉角


if __name__ == "__main__2":
    arm = Project_Arm()
    # Define start and end poses
    start_pose = SE3(0, 0, 0) * SE3.RPY(0, 0, 0)
    end_pose = SE3(200, 200, 200) * SE3.RPY(1, 0, -0.5)  # Convert mm to meters
    
    # Number of steps for interpolation
    steps = 50
    
    # Interpolate between start and end poses
    traj = start_pose.interp(end_pose, steps)
    
    # calculate inverse kinematics for each pose
    qs = [ arm.ikine_LM(traj[0]) ]
    for i in range(1,steps):
        qs.append(arm.ikine_LM(traj[i], q0 =qs[-1]))
    q = np.array([qs[i].q for i in range(steps)])
    
    # Plot the robot with pyplot and loop
    arm.plot(q, dt=0.1,loop=True,block=True)
    
    
    